{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from PIL import Image\n",
    "import time\n",
    "\n",
    "VGG_MEAN = [103.939, 116.779, 123.68]\n",
    "\n",
    "class VGGNet:\n",
    "    \"\"\"Builds VGG-16 net structure,\n",
    "       load parameters from pre-train models.\n",
    "    \"\"\"\n",
    "    def __init__(self, data_dict):\n",
    "        self.data_dict = data_dict\n",
    "    \n",
    "    def get_conv_filter(self, name):\n",
    "        return tf.constant(self.data_dict[name][0], name='conv')\n",
    "    \n",
    "    def get_fc_weight(self, name):\n",
    "        return tf.constant(self.data_dict[name][0], name='fc')\n",
    "    \n",
    "    def get_bias(self, name):\n",
    "        return tf.constant(self.data_dict[name][1], name='bias')\n",
    "    \n",
    "    def conv_layer(self, x, name):\n",
    "        \"\"\"Builds convolution layer.\"\"\"\n",
    "        with tf.name_scope(name):\n",
    "            conv_w = self.get_conv_filter(name)\n",
    "            conv_b = self.get_bias(name)\n",
    "            h = tf.nn.conv2d(x, conv_w, [1,1,1,1], padding='SAME')\n",
    "            h = tf.nn.bias_add(h, conv_b)\n",
    "            h = tf.nn.relu(h)\n",
    "            return h\n",
    "    \n",
    "    \n",
    "    def pooling_layer(self, x, name):\n",
    "        \"\"\"Builds pooling layer.\"\"\"\n",
    "        return tf.nn.max_pool(x,\n",
    "                              ksize = [1,2,2,1],\n",
    "                              strides = [1,2,2,1],\n",
    "                              padding = 'SAME',\n",
    "                              name = name)\n",
    "    \n",
    "    def fc_layer(self, x, name, activation=tf.nn.relu):\n",
    "        \"\"\"Builds fully-connected layer.\"\"\"\n",
    "        with tf.name_scope(name):\n",
    "            fc_w = self.get_fc_weight(name)\n",
    "            fc_b = self.get_bias(name)\n",
    "            h = tf.matmul(x, fc_w)\n",
    "            h = tf.nn.bias_add(h, fc_b)\n",
    "            if activation is None:\n",
    "                return h\n",
    "            else:\n",
    "                return activation(h)\n",
    "    \n",
    "    def flatten_layer(self, x, name):\n",
    "        \"\"\"Builds flatten layer.\"\"\"\n",
    "        with tf.name_scope(name):\n",
    "            # [batch_size, image_width, image_height, channel]\n",
    "            x_shape = x.get_shape().as_list()\n",
    "            dim = 1\n",
    "            for d in x_shape[1:]:\n",
    "                dim *= d\n",
    "            x = tf.reshape(x, [-1, dim])\n",
    "            return x\n",
    "    \n",
    "    def build(self, x_rgb):\n",
    "        \"\"\"Build VGG16 network structure.\n",
    "        Parameters:\n",
    "        - x_rgb: [1, 224, 224, 3]\n",
    "        \"\"\"\n",
    "        \n",
    "        start_time = time.time()\n",
    "        print 'building model ...'\n",
    "        \n",
    "        r, g, b = tf.split(x_rgb, [1,1,1], axis=3)\n",
    "        x_bgr = tf.concat(\n",
    "            [b - VGG_MEAN[0],\n",
    "             g - VGG_MEAN[1],\n",
    "             r - VGG_MEAN[2]],\n",
    "            axis = 3)\n",
    "        \n",
    "        assert x_bgr.get_shape().as_list()[1:] == [224, 224, 3]\n",
    "        \n",
    "        self.conv1_1 = self.conv_layer(x_bgr, 'conv1_1')\n",
    "        self.conv1_2 = self.conv_layer(self.conv1_1, 'conv1_2')\n",
    "        self.pool1 = self.pooling_layer(self.conv1_2, 'pool1')\n",
    "        \n",
    "        self.conv2_1 = self.conv_layer(self.pool1, 'conv2_1')\n",
    "        self.conv2_2 = self.conv_layer(self.conv2_1, 'conv2_2')\n",
    "        self.pool2 = self.pooling_layer(self.conv2_2, 'pool2')\n",
    "        \n",
    "        self.conv3_1 = self.conv_layer(self.pool2, 'conv3_1')\n",
    "        self.conv3_2 = self.conv_layer(self.conv3_1, 'conv3_2')\n",
    "        self.conv3_3 = self.conv_layer(self.conv3_2, 'conv3_3')\n",
    "        self.pool3 = self.pooling_layer(self.conv3_3, 'pool3')\n",
    "        \n",
    "        self.conv4_1 = self.conv_layer(self.pool3, 'conv4_1')\n",
    "        self.conv4_2 = self.conv_layer(self.conv4_1, 'conv4_2')\n",
    "        self.conv4_3 = self.conv_layer(self.conv4_2, 'conv4_3')\n",
    "        self.pool4 = self.pooling_layer(self.conv4_3, 'pool4')\n",
    "        \n",
    "        self.conv5_1 = self.conv_layer(self.pool4, 'conv5_1')\n",
    "        self.conv5_2 = self.conv_layer(self.conv5_1, 'conv5_2')\n",
    "        self.conv5_3 = self.conv_layer(self.conv5_2, 'conv5_3')\n",
    "        self.pool5 = self.pooling_layer(self.conv5_3, 'pool5')\n",
    "        \n",
    "        '''\n",
    "        self.flatten5 = self.flatten_layer(self.pool5, 'flatten')\n",
    "        self.fc6 = self.fc_layer(self.flatten5, 'fc6')\n",
    "        self.fc7 = self.fc_layer(self.fc6, 'fc7')\n",
    "        self.fc8 = self.fc_layer(self.fc7, 'fc8', activation=None)\n",
    "        self.prob = tf.nn.softmax(self.fc8, name='prob')\n",
    "        '''\n",
    "        \n",
    "        print 'building model finished: %4ds' % (time.time() - start_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vgg16_npy_path = '../../style_transfer_data/vgg16.npy'\n",
    "content_img_path = '../../style_transfer_data/gugong.jpg'\n",
    "style_img_path = '../../style_transfer_data/xingkong.jpeg'\n",
    "\n",
    "num_steps = 100\n",
    "learning_rate = 10\n",
    "\n",
    "lambda_c = 0.1\n",
    "lambda_s = 500\n",
    "\n",
    "output_dir = './run_style_transfer'\n",
    "\n",
    "if not os.path.exists(output_dir):\n",
    "    os.mkdir(output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def initial_result(shape, mean, stddev):\n",
    "    initial = tf.truncated_normal(shape, mean = mean, stddev = stddev)\n",
    "    return tf.Variable(initial)\n",
    "\n",
    "def read_img(img_name):\n",
    "    img = Image.open(img_name)\n",
    "    np_img = np.array(img) # (224, 224, 3)\n",
    "    np_img = np.asarray([np_img], dtype=np.int32) # (1, 224, 224, 3)\n",
    "    return np_img\n",
    "\n",
    "def gram_matrix(x):\n",
    "    \"\"\"Calulates gram matrix\n",
    "    Args:\n",
    "    - x: feaures extracted from VGG Net. shape: [1, width, height, ch]\n",
    "    \"\"\"\n",
    "    b, w, h, ch = x.get_shape().as_list()\n",
    "    features = tf.reshape(x, [b, h*w, ch]) # [ch, ch] -> (i, j)\n",
    "    # [h*w, ch] matrix -> [ch, h*w] * [h*w, ch] -> [ch, ch]\n",
    "    gram = tf.matmul(features, features, adjoint_a=True) / tf.constant(ch * w * h, tf.float32)\n",
    "    return gram\n",
    "    \n",
    "\n",
    "result = initial_result((1, 224, 224, 3), 127.5, 20)\n",
    "\n",
    "content_val = read_img(content_img_path)\n",
    "style_val = read_img(style_img_path)\n",
    "\n",
    "content = tf.placeholder(tf.float32, shape=[1, 224, 224, 3])\n",
    "style = tf.placeholder(tf.float32, shape=[1, 224, 224, 3])\n",
    "\n",
    "data_dict = np.load(vgg16_npy_path).item()\n",
    "vgg_for_content = VGGNet(data_dict)\n",
    "vgg_for_style = VGGNet(data_dict)\n",
    "vgg_for_result = VGGNet(data_dict)\n",
    "\n",
    "vgg_for_content.build(content)\n",
    "vgg_for_style.build(style)\n",
    "vgg_for_result.build(result)\n",
    "\n",
    "content_features = [\n",
    "    vgg_for_content.conv1_2,\n",
    "    # vgg_for_content.conv2_2,\n",
    "    # vgg_for_content.conv3_3,\n",
    "    # vgg_for_content.conv4_3,\n",
    "    # vgg_for_content.conv5_3\n",
    "]\n",
    "\n",
    "result_content_features = [\n",
    "    vgg_for_result.conv1_2,\n",
    "    # vgg_for_result.conv2_2,\n",
    "    # vgg_for_result.conv3_3,\n",
    "    # vgg_for_result.conv4_3,\n",
    "    # vgg_for_result.conv5_3\n",
    "]\n",
    "\n",
    "# feature_size, [1, width, height, channel]\n",
    "style_features = [\n",
    "    # vgg_for_style.conv1_2,\n",
    "    # vgg_for_style.conv2_2,\n",
    "    # vgg_for_style.conv3_3,\n",
    "    vgg_for_style.conv4_3,\n",
    "    # vgg_for_style.conv5_3\n",
    "]\n",
    "style_gram = [gram_matrix(feature) for feature in style_features]\n",
    "\n",
    "result_style_features = [\n",
    "    # vgg_for_result.conv1_2,\n",
    "    # vgg_for_result.conv2_2,\n",
    "    # vgg_for_result.conv3_3,\n",
    "    vgg_for_result.conv4_3,\n",
    "    # vgg_for_result.conv5_3\n",
    "]\n",
    "result_style_gram = \\\n",
    "    [gram_matrix(feature) for feature in result_style_features]\n",
    "\n",
    "content_loss = tf.zeros(1, tf.float32)\n",
    "# zip: [1, 2], [3, 4], zip([1,2], [3,4]) -> [(1, 3), (2, 4)]\n",
    "# shape: [1, width, height, channel]\n",
    "for c, c_ in zip(content_features, result_content_features):\n",
    "    content_loss += tf.reduce_mean((c - c_) ** 2, [1, 2, 3])\n",
    "    \n",
    "style_loss = tf.zeros(1, tf.float32)\n",
    "for s, s_ in zip(style_gram, result_style_gram):\n",
    "    style_loss += tf.reduce_mean((s - s_) ** 2, [1, 2])\n",
    "\n",
    "loss = content_loss * lambda_c + style_loss * lambda_s\n",
    "train_op = tf.train.AdamOptimizer(learning_rate).minimize(loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_op = tf.global_variables_initializer()\n",
    "\n",
    "with tf.Session() as sess:\n",
    "    sess.run(init_op)\n",
    "    for step in range(num_steps):\n",
    "        loss_value, content_loss_value, style_loss_value, _ \\\n",
    "            = sess.run([loss, content_loss, style_loss, train_op],\n",
    "                     feed_dict = {\n",
    "                         content: content_val,\n",
    "                         style: style_val,\n",
    "                     })\n",
    "        print 'step: %d, loss_value: %8.4f, content_loss: %8.4f, style_loss: %8.4f' \\\n",
    "            % (step+1,\n",
    "               loss_value[0],\n",
    "               content_loss_value[0],\n",
    "               style_loss_value[0])\n",
    "        result_img_path = os.path.join(\n",
    "            output_dir, 'result-%05d.jpg' % (step+1))\n",
    "        result_val = result.eval(sess)[0]\n",
    "        result_val = np.clip(result_val, 0, 255)\n",
    "        img_arr = np.asarray(result_val, np.uint8)\n",
    "        img = Image.fromarray(img_arr)\n",
    "        img.save(result_img_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
